#include "abc_archive.h"
#include <unordered_set>

namespace abc
{

    Path parsePath(const std::string &moduleName, const std::string &nsName = "")
    {
        Path path;
        path.originModulePath = moduleName;
        std::vector<std::string> extracted;
        if (validateAndExtractPath(moduleName, extracted) && extracted.size() == 2)
        {
            path.module = extracted[0];
            path.package = extracted[1];
        }
        else
        {
            path.package = moduleName;
        }
        if (nsName != "")
        {
            path.namespaces.emplace_back(nsName);
        }
        return path;
    }

    void packageFunction(AbcNodePtr node, AbckitCoreFunction *function)
    {
        auto func = std::make_shared<AbcFunction>();
        func->name = abckitApi->string.getFunctionName(function);
        func->coreFunction = function;
        func->path = node->path;
        func->parentNode = node;
        if (abckitApi->node.isModuleNode(node->coreNode) || abckitApi->node.isNamespaceNode(node->coreNode))
        {
            func->isTopLevel = true;
        }
        func->isStatic = abckitApi->function.functionIsStatic(function);
        node->functions.emplace_back(func);
    }

    void collectFunctions(AbcNodePtr abcNode)
    {
        if (auto *module = getAs<AbckitCoreModule>(abcNode->coreNode))
        {
            abckitApi->node.moduleEnumerateTopLevelFunctions(module, abcNode.get(), [](AbckitCoreFunction *function, void *data)
                                                             {
                AbcNode *node = static_cast<AbcNode *>(data);
                auto nodePtr = AbcNodePtr(node, [](AbcNode*){});
                packageFunction(nodePtr,function);
                return true; });

            return;
        }
        if (auto *ns = getAs<AbckitCoreNamespace>(abcNode->coreNode))
        {
            abckitApi->node.namespaceEnumerateTopLevelFunctions(ns, abcNode.get(), [](AbckitCoreFunction *function, void *data)
                                                                {
                AbcNode *node = static_cast<AbcNode *>(data);
                auto nodePtr = AbcNodePtr(node, [](AbcNode*){});
                packageFunction(nodePtr,function);
                return true; });
            return;
        }
        if (auto *klass = getAs<AbckitCoreClass>(abcNode->coreNode))
        {
            abckitApi->node.classEnumerateMethods(klass, abcNode.get(), [](AbckitCoreFunction *method, void *data)
                                                  {
                AbcNode *node = static_cast<AbcNode *>(data);
                auto nodePtr = AbcNodePtr(node, [](AbcNode*){});
                packageFunction(nodePtr,method);
                return true; });
            return;
        }
    }

    struct NodeContext
    {
        AbcArchive *archive;
        AbcNodePtr parentNode;
    };

    AbcNodePtr packeNode(CoreNode coreNode, AbcNodePtr parentNode = nullptr)
    {
        auto node = AbcNodePtr(new AbcNode());
        node->type = abckitApi->node.getNodeType(coreNode);
        node->name = abckitApi->string.getNodeName(coreNode);
        node->path = parsePath(node->type == AbcNodeType::MODULE ? node->name : parentNode->path.originModulePath, (node->type == AbcNodeType::NAMESPACE) ? node->name : "");
        node->coreNode = coreNode;
        if (parentNode)
        {
            node->parentNode = parentNode;
        }
        collectFunctions(node);
        return node;
    }

    void AbcArchiveFactory::collectNode(AbcArchive *archive, CoreNode coreNode, AbcNodePtr parentNode = nullptr)
    {
        auto node = packeNode(coreNode, parentNode);
        archive->_nodes.emplace_back(node);
        NodeContext context;
        context.archive = archive;
        context.parentNode = node;
        if (auto *module = getAs<AbckitCoreModule>(coreNode))
        {
            abckitApi->node.moduleEnumerateNamespaces(module, &context, [](AbckitCoreNamespace *coreNode, void *data)
                                                      {
                NodeContext *context = reinterpret_cast<NodeContext *>(data);
                AbcArchiveFactory::collectNode(context->archive,coreNode,context->parentNode);
                return true; });
            abckitApi->node.moduleEnumerateClasses(module, &context, [](AbckitCoreClass *coreNode, void *data)
                                                   {
                NodeContext *context = reinterpret_cast<NodeContext *>(data);
                AbcArchiveFactory::collectNode(context->archive,coreNode,context->parentNode);
                return true; });
        }
        if (auto *ns = getAs<AbckitCoreNamespace>(coreNode))
        {
            abckitApi->node.namespaceEnumerateNamespaces(ns, &context, [](AbckitCoreNamespace *coreNode, void *data)
                                                         {
                NodeContext *context = reinterpret_cast<NodeContext *>(data);
                AbcArchiveFactory::collectNode(context->archive,coreNode,context->parentNode);
                return true; });
            abckitApi->node.namespaceEnumerateClasses(ns, &context, [](AbckitCoreClass *coreNode, void *data)
                                                      {
                NodeContext *context = reinterpret_cast<NodeContext *>(data);
                AbcArchiveFactory::collectNode(context->archive,coreNode,context->parentNode);
                return true; });
        }
    }

    void AbcArchiveFactory::cleanNode(AbcArchive *archive)
    {
        // 1. 收集所有被引用的父节点（存活的父节点）
        std::unordered_set<AbcNode *> referencedParents;
        for (const auto &node : archive->_nodes)
        {
            if (auto parent = node->parentNode)
            {
                referencedParents.insert(parent.get());
            }
        }

        // 2. 移除满足条件的节点
        archive->_nodes.erase(
            std::remove_if(
                archive->_nodes.begin(),
                archive->_nodes.end(),
                [&referencedParents](const AbcNodePtr &node)
                {
                    // 条件1：functions 为空
                    bool hasNoFunctions = node->functions.empty();
                    // 条件2：未被其他节点引用（不在 referencedParents 中）
                    bool isOrphaned = referencedParents.find(node.get()) == referencedParents.end();
                    return hasNoFunctions && isOrphaned;
                }),
            archive->_nodes.end());
    }

    AbcFuncType AbcArchiveFactory::identifyFunctionType(const AbcFunctionPtr &function, bool isPage = false)
    {
        auto funcName = function->name;
        if (funcName == "func_main_0")
        {
            return AbcFuncType::MAIN;
        }
        if (abckitApi->function.functionIsAnonymous(function->coreFunction))
        {
            return AbcFuncType::LAMBDA;
        }

        if (auto parentNode = function->parentNode)
        {
            if (abckitApi->node.isClassNode(parentNode->coreNode) && parentNode->name == funcName)
            {
                return AbcFuncType::CONSTRUCTOR;
            }
        }

        if (!function->isStatic && funcName == "initialRender")
        {
            return AbcFuncType::BUILD;
        }

        if (isPage && (funcName == "rerender" || funcName == "updateStateVars" || funcName == "setInitiallyProvidedValue" || funcName == "purgeVariableDependenciesOnElmtId" || funcName == "getEntryName"))
        {
            return AbcFuncType::OTHER;
        }

        if (funcName == "aboutToAppear" || funcName == "aboutToBeDeleted" || funcName == "onBackup" || funcName == "onRestore" || funcName == "onBackground" || funcName == "onCreate" || funcName == "onDestroy" || funcName == "onForeground" || funcName == "onWindowStageCreate" || funcName == "onWindowStageDestroy")
        {
            return AbcFuncType::LIFECYCLE;
        }

        return AbcFuncType::NORMAL;
    }

    void identifyGetOrSetFunctionType(const abc::AbcNodePtr &node, std::unordered_set<std::string> attrNames)
    {
        for (const auto &func : node->functions)
        {
            if (func->type != AbcFuncType::NORMAL)
            {
                continue;
            }
            auto funcName = func->name;
            if (attrNames.contains(funcName))
            {
                func->type = AbcFuncType::GETTER;
                auto it = std::find_if(node->functions.begin(), node->functions.end(), [funcName](auto func)
                                       { return func->name == funcName + "^1"; });
                if (it != node->functions.end())
                {
                    (*it)->type = AbcFuncType::SETTER;
                }
            }
        }
    }

    void AbcArchiveFactory::identifyFunctionType(AbcArchive *archive)
    {
        const std::string suffix = "^1";
        for (const auto &node : archive->_nodes)
        {
            // struct Page页面中有一些方法是预定了的
            bool isPage = std::any_of(node->functions.begin(), node->functions.end(),
                                      [](const AbcFunctionPtr &func)
                                      { return func->name == "initialRender"; });
            // 第一轮先分析基础的类型
            std::unordered_set<std::string> attrNames;
            for (const auto &func : node->functions)
            {
                if (func->name.size() >= suffix.size() && func->name.compare(func->name.size() - suffix.size(), suffix.size(), suffix) == 0)
                {
                    std::string methodName = func->name.substr(0, func->name.size() - suffix.size());
                    attrNames.insert(methodName);
                }
                func->type = AbcArchiveFactory::identifyFunctionType(func, isPage); // 识别并赋值类型
            }

            // 第二轮对属性Get、Set方法进行标注
            identifyGetOrSetFunctionType(node, attrNames);
        }
    }

    struct FuncContext
    {
        std::set<AbckitCoreFunction *> lambdas;
    };

    void getLambdaFunction(AbckitCoreFunction *lamdaFunction, const AbcFunctionPtr &parentFunc, const AbcNodePtr &node)
    {
        auto tmpNode = node;
        while (tmpNode)
        {
            auto it = std::find_if(
                tmpNode->functions.begin(),
                tmpNode->functions.end(),
                [lamdaFunction](const AbcFunctionPtr &func)
                {
                    return func->coreFunction == lamdaFunction;
                });
            if (it != tmpNode->functions.end())
            {
                parentFunc->lambdas.insert(*it);
                return;
            }
            tmpNode = tmpNode->parentNode;
        }
    }

    void collectDirectLambdas(FuncContext *funcContext, const AbcFunctionPtr &func, const AbcNodePtr &node)
    {
        if (funcContext->lambdas.size() < 1)
        {
            return;
        }

        for (const auto &abckitCoreFunction : funcContext->lambdas)
        {
            getLambdaFunction(abckitCoreFunction, func, node);
        }
    }

    void collectDirectLambdas(const AbcFunctionPtr &func, const AbcNodePtr &node)
    {
        auto graph = abckitApi->graph.createGraphFromFunction(func->coreFunction);

        FuncContext funcContext;
        abckitApi->graph.gVisitBlocksRpo(graph, &funcContext, [](AbckitBasicBlock *basicBlock, void *data)
                                         {
            FuncContext *context = reinterpret_cast<FuncContext *>(data);
            auto inst = abckitApi->inst.bbGetFirstInst(basicBlock);
            while(inst != nullptr){
                if(abckitApi->inst.iGetOpcode(inst) == AbckitIsaApiDynamicOpcode::ABCKIT_ISA_API_DYNAMIC_OPCODE_DEFINEFUNC){
                    auto method = abckitApi->inst.iGetFunction(inst);
                    context->lambdas.insert(method);
                }

                inst = abckitApi->inst.iGetNext(inst);
            }
            return true; });
        collectDirectLambdas(&funcContext, func, node);
        abckitApi->graph.destroyGraph(graph);
    }

    void collectAllChildLambdas(std::set<AbcFunctionPtr> &children, const AbcFunctionPtr &func)
    {
        children.insert(func);
        for (const auto &lambdaFunc : func->lambdas)
        {
            collectAllChildLambdas(children, lambdaFunc);
        }
    }

    void collectAllChildLambdas(const AbcFunctionPtr &func)
    {
        std::set<AbcFunctionPtr> children;
        for (const auto &lambdaFunc : func->lambdas)
        {
            collectAllChildLambdas(children, lambdaFunc);
        }
        for (const auto &lambdaFunc : children)
        {
            func->lambdas.insert(lambdaFunc);
        }
    }

    /**
     * 遍历方法体收集lambda函数
     */
    void AbcArchiveFactory::scanFunctionsForLambdas(AbcArchive *archive)
    {
        for (const auto &node : archive->_nodes)
        {
            for (const auto &func : node->functions)
            {
                collectDirectLambdas(func, node);
            }

            for (const auto &func : node->functions)
            {
                collectAllChildLambdas(func);
            }
        }
    }

    AbcArchive::Ptr AbcArchiveFactory::create(const std::string &path)
    {
        if (path.empty())
            return nullptr;
        auto *archive = new AbcArchive(path);
        archive->rawFile_ = abckitApi->io.openAbc(archive->filePath_);
        abckitApi->node.fileEnumerateModules(archive->rawFile_, archive, [](AbckitCoreModule *module, void *data)
                                             {
            AbcArchive *archive = reinterpret_cast<AbcArchive *>(data);
            AbcArchiveFactory::collectNode(archive,module);
            
            return true; });
        AbcArchiveFactory::cleanNode(archive);
        AbcArchiveFactory::identifyFunctionType(archive);
        AbcArchiveFactory::scanFunctionsForLambdas(archive);
        return AbcArchive::Ptr(archive);
    }
}